from fastapi import WebSocket
from typing import Dict


class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str,WebSocket] = {}

    async def connect(self, username:str, websocket:WebSocket):
        await websocket.accept()
        self.active_connections[username] = websocket
        await self.broadcast(f"用户{username}加入了聊天室")

    async def disconnect(self, username:str):
        if username in self.active_connections:
            del self.active_connections[username]

    async def send_personal_message(self, message:str, to_user: str):
        if to_user in self.active_connections:
            await self.active_connections[to_user].send(message)

    async def broadcast(self, message:str):
        for ws in self.active_connections.values():
            await ws.send_text(message)
